#!/usr/bin/python
# Global modules.
import subprocess
import sys
import os
import time

import fasta_read

# usage
if len(sys.argv) == 1:
	print 
	print "usage: -fasta file -agp file -out file [-id text] [-addgap]"
	print "-fasta: an assembly."
	print "-agp: an agp for said assembly."
	print "-out: fasta file out."
	print "-id: scaffold to linearlize."
	print "-addgap: if there are no gaps in agp."
	print """Linearizes scaffolds.
	"""
	print
	sys.exit()
	
# arguments.
infasta = ""
inagp = ""
scafid = ""
outfile = ""
addgap = False
for i in range(len(sys.argv)):
	if sys.argv[i] == "-fasta":
		infasta = sys.argv[i+1]
	if sys.argv[i] == "-agp":
		inagp = sys.argv[i+1]	
	if sys.argv[i] == "-id":
		scafid = sys.argv[i+1]	
	if sys.argv[i] == "-out":
		outfile = sys.argv[i+1]			
	if sys.argv[i] == "-addgap":
		addgap = True
	

# Read in fasta.
data = fasta_read.read_fasta_dict(infasta)
	
# Read in scaffolds.
pstop = -1
prevscaf = ""
seq = ""
fin = open(inagp,"r")
fout = open(outfile,"w")
for line in fin:
	tmp = line.split("\t")
	
	# Skip till find scaffold.
	if scafid != "":
		if tmp[0] != scafid: continue
	
	if tmp[0] != prevscaf:
		if prevscaf != "":
			fout.write(">%s\n" % prevscaf)
			fout.write("%s\n" % seq)
		seq = ""
		prevscaf = tmp[0]
	
	# add gap.
	if addgap == True:
		if pstop != -1:
			cstart = int(tmp[1])
			seq += "N" * (cstart - pstop)
	
	# Add contig.
	if tmp[4] == "W":
		# add seq.
		seq += data[tmp[5]]
		
		# increment previous stop.
		pstop = int(tmp[2])
		continue
		
	# Add gap.
	if addgap == False:
		nn = ""
		if tmp[4] == "N":
			sz = int(tmp[5])
			for i in range(sz):
				nn += "N"
			seq += nn
			continue
fin.close()
fout.close()
